import os
import numpy as np
import torch
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans, MiniBatchKMeans
from tqdm import tqdm
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

from utils.common import check_and_create_directory, list_all_files, chunk_list

if __name__ == '__main__':
    # todo 提取所有的数据，然后合并在一起，看看大小
    input_path = '/home/zry/datasets/building/train/rgbs'
    output_path = '/home/zry/datasets/building/train/seg'
    tmp_path = '/home/zry/datasets/building/train/tmp'
    batch_size = 64
    n_clusters = 4  # 假设我们选择5个聚类

    kmeans = MiniBatchKMeans(n_clusters=n_clusters, batch_size=batch_size, n_init="auto", reassignment_ratio=0)
    
    # 所有pth
    paths = list_all_files(tmp_path, '.pth')
    path_batch = chunk_list(paths, batch_size)
    
    for batch in path_batch:
        encs = []
        for i, path in tqdm(enumerate(batch), total=len(batch)):
            f = torch.load(path, "cpu")
            idx, seg_images, enc = f['idx'], f['seg_images'], f['enc']
            encs.append(enc)
        
        encs = torch.cat(encs)
        encs = encs.detach().numpy()

        kmeans = kmeans.partial_fit(encs)
    
    # 使用K-Means进行聚类
    labels = kmeans.predict(encs)
    
    # 使用TSNE进行降维
    tsne = TSNE(n_components=2, initialization="random", n_jobs=8)
    tsne_results = []
    for batch in path_batch:
        encs = []
        for i, path in tqdm(enumerate(batch), total=len(batch)):
            f = torch.load(path, "cpu")
            idx, seg_images, enc = f['idx'], f['seg_images'], f['enc']
            encs.append(enc)
        
        encs = torch.cat(encs)
        encs = encs.detach().numpy()
        tsne_results.append(tsne.fit(encs))

    tsne_results = np.vstack(tsne_results)

    # 可视化TSNE结果并根据聚类结果着色
    plt.figure(figsize=(100, 60))
    scatter = plt.scatter(tsne_results[:, 0], tsne_results[:, 1], c=labels, marker='o', cmap='viridis')
    plt.colorbar(scatter)
    plt.grid(True)
    
    # 保存图片文件
    plt.savefig('tsne_reduction_with_clusters.png')
    plt.show()
